#include "imm.h"



// 创建IMM
Imm* create_imm(KalmanFilter** models, Matrix*** model_trans, Matrix* P_trans, Matrix* U_prob, int mode_cnt) {
    Imm* imm = (Imm*)malloc(sizeof(Imm));
    imm->mode_cnt = mode_cnt;
    imm->models = (KalmanFilter**)malloc(mode_cnt * sizeof(KalmanFilter*));
    for (int i = 0; i < mode_cnt; i++) {
        imm->models[i] = models[i];
    }
    
    // 分配模型转换矩阵内存
    imm->model_trans = (Matrix***)malloc(mode_cnt * sizeof(Matrix**));
    for (int i = 0; i < mode_cnt; i++) {
        imm->model_trans[i] = (Matrix**)malloc(mode_cnt * sizeof(Matrix*));
        for (int j = 0; j < mode_cnt; j++) {
            imm->model_trans[i][j] = matrix_copy(model_trans[i][j]);
        }
    }
    
    imm->P_trans = matrix_copy(P_trans);
    imm->U_prob = matrix_copy(U_prob);
    
    return imm;
}

// 释放IMM
void free_imm(Imm* imm) {
    if (imm) {
        for (int i = 0; i < imm->mode_cnt; i++) {
            free_kalman_filter(imm->models[i]);
            for (int j = 0; j < imm->mode_cnt; j++) {
                free_matrix(imm->model_trans[i][j]);
            }
            free(imm->model_trans[i]);
        }
        free(imm->models);
        free(imm->model_trans);
        free_matrix(imm->P_trans);
        free_matrix(imm->U_prob);
        free(imm);
    }
}

// IMM滤波
Matrix* imm_filt(Imm* imm, Matrix* Z) {
    int mode_cnt = imm->mode_cnt;
    
    // Step 1: 输入混合
    Matrix* u = matrix_multiply(matrix_transpose(imm->P_trans), imm->U_prob);
    
    Matrix* mu = create_matrix(mode_cnt, mode_cnt);
    for (int i = 0; i < mode_cnt; i++) {
        for (int j = 0; j < mode_cnt; j++) {
            mu->data[i][j] = imm->P_trans->data[i][j] * imm->U_prob->data[i][0] / u->data[j][0];
        }
    }
    printf("===  %d imm filter ===\n");
    // 混合状态
    Matrix** X_mix = (Matrix**)malloc(mode_cnt * sizeof(Matrix*));
    for (int j = 0; j < mode_cnt; j++) {
        X_mix[j] = create_matrix(imm->models[j]->X->rows, 1);
          printf("===  %d imm for ===\n");
        for (int i = 0; i < mode_cnt; i++) {
            Matrix* temp = matrix_multiply(imm->model_trans[j][i], imm->models[i]->X);
            Matrix* temp2 = matrix_scalar_multiply(temp, mu->data[i][j]);
            Matrix* temp3 = matrix_add(X_mix[j], temp2);
            memcpy(X_mix[j]->data, temp3->data, X_mix[j]->rows * X_mix[j]->cols * sizeof(double));
            
            free_matrix(temp);
            free_matrix(temp2);
            free_matrix(temp3);
            printf("===  %d imm for in ===\n");
        }
        
    }
     printf("===  %d imm end ===\n");
    // 混合协方差
    Matrix** P_mix = (Matrix**)malloc(mode_cnt * sizeof(Matrix*));
    for (int j = 0; j < mode_cnt; j++) {
        P_mix[j] = create_matrix(imm->models[j]->P->rows, imm->models[j]->P->cols);
        for (int i = 0; i < mode_cnt; i++) {
            Matrix* temp = matrix_subtract(imm->models[i]->X, X_mix[i]);
            Matrix* temp2 = matrix_multiply(temp, matrix_transpose(temp));
            Matrix* P = matrix_add(imm->models[i]->P, temp2);
            free_matrix(temp);
            free_matrix(temp2);
            
            Matrix* temp3 = matrix_multiply(imm->model_trans[j][i], P);
            Matrix* temp4 = matrix_multiply(temp3, matrix_transpose(imm->model_trans[j][i]));
            Matrix* temp5 = matrix_scalar_multiply(temp4, mu->data[i][j]);
            Matrix* temp6 = matrix_add(P_mix[j], temp5);
            memcpy(P_mix[j]->data, temp6->data, P_mix[j]->rows * P_mix[j]->cols * sizeof(double));
            free_matrix(P);
            free_matrix(temp3);
            free_matrix(temp4);
            free_matrix(temp5);
            free_matrix(temp6);
        }
    }
    
    // Step 2: 滤波
    for (int j = 0; j < mode_cnt; j++) {
        memcpy(imm->models[j]->X->data, X_mix[j]->data, imm->models[j]->X->rows * imm->models[j]->X->cols * sizeof(double));
        memcpy(imm->models[j]->P->data, P_mix[j]->data, imm->models[j]->P->rows * imm->models[j]->P->cols * sizeof(double));
        kalman_filt(imm->models[j], Z);
        free_matrix(X_mix[j]);
        free_matrix(P_mix[j]);
    }
    free(X_mix);
    free(P_mix);
    
    // Step 3: 更新概率
    double total_lambda = 0.0;
    double* lambda = (double*)malloc(mode_cnt * sizeof(double));
    
    for (int j = 0; j < mode_cnt; j++) {
        KalmanFilter* mode = imm->models[j];
        Matrix* temp = matrix_multiply(mode->H, mode->X_pre);
        Matrix* D = matrix_subtract(Z, temp);
        free_matrix(temp);
        
        // 计算S并确保对称
        temp = matrix_multiply(mode->H, mode->P_pre);
        Matrix* temp2 = matrix_transpose(mode->H);
        Matrix* temp3 = matrix_multiply(temp, temp2);
        free_matrix(temp);
        free_matrix(temp2);
        
        Matrix* S = matrix_add(temp3, mode->R);
        free_matrix(temp3);
        
        // 强制对称
        for (int i = 0; i < S->rows; i++) {
            for (int k = i + 1; k < S->cols; k++) {
                S->data[i][k] = S->data[k][i] = (S->data[i][k] + S->data[k][i]) * 0.5;
            }
        }
        
        // 添加正则化
        Matrix* reg = matrix_scalar_multiply(matrix_eye(S->rows), REG_EPS);
        Matrix* S_reg = matrix_add(S, reg);
        free_matrix(reg);
        
        Matrix* S_inv = create_matrix(S_reg->rows, S_reg->cols);
        int inv_success = matrix_inverse(S_reg, S_inv);
        
        double mahalanobis = 0.0;
        if (inv_success) {
            Matrix* temp4 = matrix_multiply(matrix_transpose(D), S_inv);
            Matrix* temp5 = matrix_multiply(temp4, D);
            mahalanobis = temp5->data[0][0];
            free_matrix(temp4);
            free_matrix(temp5);
        }
        
        double det_S = matrix_det(S_reg);
        free_matrix(S_reg);
        
        double Lambda;
        if (det_S <= 0 || !inv_success) {
            Lambda = MIN_LAMBDA;
        } else {
            int k = D->rows;
            double norm_const = 1.0 / (pow(2 * PI, k / 2.0) * sqrt(det_S));
            Lambda = norm_const * exp(-0.5 * mahalanobis);
            if (Lambda < MIN_LAMBDA) Lambda = MIN_LAMBDA;
        }
        
        lambda[j] = Lambda * u->data[j][0];
        total_lambda += lambda[j];
        
        free_matrix(D);
        free_matrix(S);
        free_matrix(S_inv);
    }
    
    // 归一化
    if (total_lambda <= 0) {
        for (int j = 0; j < mode_cnt; j++) {
            imm->U_prob->data[j][0] = 1.0 / mode_cnt;
        }
    } else {
        for (int j = 0; j < mode_cnt; j++) {
            imm->U_prob->data[j][0] = lambda[j] / total_lambda;
        }
    }
    
    free(lambda);
    free_matrix(u);
    return matrix_copy(imm->U_prob);
}





// 匀速模型(CV)
KalmanFilter* kf_cv() {
    Matrix* A = create_matrix(6, 6);
    A->data[0][0] = 1; A->data[0][1] = DT;
    A->data[1][1] = 1;
    A->data[2][2] = 1; A->data[2][3] = DT;
    A->data[3][3] = 1;
    A->data[4][4] = 1; A->data[4][5] = DT;
    A->data[5][5] = 1;
    
    Matrix* H = matrix_eye(6);
    
    KalmanFilter* kf = create_kalman_filter(A, H);
    
    // 设置观测噪声
    double r[] = {5.0, 2.0, 5.0, 2.0, 5.0, 2.0};
    for (int i = 0; i < 6; i++) {
        kf->R->data[i][i] = r[i];
    }
    
    free_matrix(A);
    free_matrix(H);
    return kf;
}

// CT_xy模型
KalmanFilter* kf_ct_xy() {
    double dtheta = PI / 180 * 5;
    double theta = dtheta * DT;
    
    Matrix* A = create_matrix(7, 7);
    A->data[0][0] = 1;
    A->data[0][1] = sin(theta) / dtheta;
    A->data[0][3] = -(1 - cos(theta)) / dtheta;
    A->data[1][1] = cos(theta);
    A->data[1][3] = -sin(theta);
    A->data[2][1] = (1 - cos(theta)) / dtheta;
    A->data[2][2] = 1;
    A->data[2][3] = sin(theta) / dtheta;
    A->data[3][1] = sin(theta);
    A->data[3][3] = cos(theta);
    A->data[4][4] = 1;
    A->data[4][5] = DT;
    A->data[5][5] = 1;
    A->data[6][6] = 1;
    
    Matrix* H = create_matrix(6, 7);
    for (int i = 0; i < 6; i++) {
        H->data[i][i] = 1;
    }
    
    KalmanFilter* kf = create_kalman_filter(A, H);
    
    // 设置观测噪声
    double r[] = {5.0, 2.0, 5.0, 2.0, 5.0, 2.0};
    for (int i = 0; i < 6; i++) {
        kf->R->data[i][i] = r[i];
    }
    
    free_matrix(A);
    free_matrix(H);
    return kf;
}

// CT_yz模型
KalmanFilter* kf_ct_yz() {
    double dtheta = PI / 180 * 5;
    double theta = dtheta * DT;
    
    Matrix* A = create_matrix(7, 7);
    A->data[0][0] = 1;
    A->data[1][1] = 1; A->data[1][2] = DT;
    A->data[2][2] = 1;
    A->data[3][3] = 1;
    A->data[3][4] = sin(theta) / dtheta;
    A->data[3][6] = -(1 - cos(theta)) / dtheta;
    A->data[4][4] = cos(theta);
    A->data[4][6] = -sin(theta);
    A->data[5][4] = (1 - cos(theta)) / dtheta;
    A->data[5][5] = 1;
    A->data[5][6] = sin(theta) / dtheta;
    A->data[6][4] = sin(theta);
    A->data[6][6] = cos(theta);
    
    Matrix* H = create_matrix(6, 7);
    for (int i = 0; i < 6; i++) {
        H->data[i][i] = 1;
    }
    
    KalmanFilter* kf = create_kalman_filter(A, H);
    
    // 设置观测噪声
    double r[] = {5.0, 2.0, 5.0, 2.0, 5.0, 2.0};
    for (int i = 0; i < 6; i++) {
        kf->R->data[i][i] = r[i];
    }
    
    free_matrix(A);
    free_matrix(H);
    return kf;
}

// CT_xz模型
KalmanFilter* kf_ct_xz() {
    double dtheta = PI / 180 * 5;
    double theta = dtheta * DT;
    
    Matrix* A = create_matrix(7, 7);
    A->data[0][0] = 1;
    A->data[0][1] = sin(theta) / dtheta;
    A->data[0][3] = -(1 - cos(theta)) / dtheta;
    A->data[1][1] = cos(theta);
    A->data[1][3] = -sin(theta);
    A->data[2][2] = 1; A->data[2][3] = DT;
    A->data[3][1] = sin(theta);
    A->data[3][3] = cos(theta);
    A->data[4][4] = 1;
    A->data[4][5] = sin(theta) / dtheta;
    A->data[4][6] = -(1 - cos(theta)) / dtheta;
    A->data[5][5] = cos(theta);
    A->data[5][6] = -sin(theta);
    A->data[6][5] = sin(theta);
    A->data[6][6] = cos(theta);
    
    Matrix* H = create_matrix(6, 7);
    for (int i = 0; i < 6; i++) {
        H->data[i][i] = 1;
    }
    
    KalmanFilter* kf = create_kalman_filter(A, H);
    
    // 设置观测噪声
    double r[] = {5.0, 2.0, 5.0, 2.0, 5.0, 2.0};
    for (int i = 0; i < 6; i++) {
        kf->R->data[i][i] = r[i];
    }
    
    free_matrix(A);
    free_matrix(H);
    return kf;
}

// 创建IMM-CVT滤波器
Imm* imm_cvt() {
    int mode_cnt = 4; // 4个模型: CV, CT_xy, CT_yz, CT_xz
    
    // 创建各个模型
    KalmanFilter* models[4];
    models[0] = kf_cv();    // 匀速模型
    models[1] = kf_ct_xy(); // CT_xy模型
    models[2] = kf_ct_yz(); // CT_yz模型
    models[3] = kf_ct_xz(); // CT_xz模型
    
    // 分配模型转换矩阵 (4x4矩阵数组)
    Matrix*** model_trans = (Matrix***)malloc(mode_cnt * sizeof(Matrix**));
    for (int i = 0; i < mode_cnt; i++) {
        model_trans[i] = (Matrix**)malloc(mode_cnt * sizeof(Matrix*));
        for (int j = 0; j < mode_cnt; j++) {
            // 确定转换矩阵的维度 (取两个模型状态维度的最大值)
            int rows = (models[j]->A->rows > models[i]->A->rows) ? 
                       models[j]->A->rows : models[i]->A->rows;
            int cols = models[i]->A->rows;
            
            model_trans[i][j] = create_matrix(rows, cols);
            
            // 初始化转换矩阵为适当维度的单位矩阵
            // 这是简化处理，实际应用中应根据模型间的物理关系设置
            for (int r = 0; r < rows && r < cols; r++) {
                model_trans[i][j]->data[r][r] = 1.0;
            }
        }
    }
    
    // 初始化转移概率矩阵 (4x4)
    Matrix* P_trans = create_matrix(mode_cnt, mode_cnt);
    // 对角线上的值较大，表示模型保持不变的概率较高
    double p = 0.8; // 自转移概率
    double q = (1 - p) / (mode_cnt - 1); // 转移到其他模型的概率
    
    for (int i = 0; i < mode_cnt; i++) {
        for (int j = 0; j < mode_cnt; j++) {
            P_trans->data[i][j] = (i == j) ? p : q;
        }
    }
    
    // 初始化模型概率 (均匀分布)
    Matrix* U_prob = create_matrix(mode_cnt, 1);
    for (int i = 0; i < mode_cnt; i++) {
        U_prob->data[i][0] = 1.0 / mode_cnt;
    }
    
    // 创建并返回IMM滤波器
    return create_imm(models, model_trans, P_trans, U_prob, mode_cnt);
}

// 打印矩阵内容（调试用）
void print_matrix(Matrix* mat) {
    if (!mat) {
        printf("矩阵为空\n");
        return;
    }
    
    for (int i = 0; i < mat->rows; i++) {
        for (int j = 0; j < mat->cols; j++) {
            printf("%10.6f ", mat->data[i][j]);
        }
        printf("\n");
    }
    printf("\n");
}

